#include "kmm.h"
#include "vmm.h"
#include "dbgio.h"
kmm_heap_t *current_heap;
kmm_heap_t address_space;
uint32_t palloc;
void kmm_setup_heap(kmm_heap_t *heap, void *base_ptr, size_t initial_size, physical_addr page){
	kmm_block_hdr_t *start_block_hdr;
	kmm_block_hdr_t *end_block_hdr;
	palloc = (uint32_t) base_ptr;
	//tprintf("[0x%X] - _kmm_setup_heap() \n",(void *)kmm_setup_heap);
	heap->start_ptr = (uint32_t) base_ptr;
	heap->free_ptr  = (uint32_t) base_ptr;
	memset(base_ptr, 0, sizeof(kmm_block_hdr_t));
	start_block_hdr = (kmm_block_hdr_t *) base_ptr;
	//tprintf("   __start_block_hdr = 0x%x\n",(uint32_t) start_block_hdr);
	//tprintf("   __initial_size    = 0x%x\n",(uint32_t) initial_size);
	if (initial_size == 0){
		start_block_hdr->next_ptr = (uint32_t) start_block_hdr;
		start_block_hdr->prev_ptr = (uint32_t) start_block_hdr;
		start_block_hdr->flags = KMM_EOM_FLAG;
	} else {		
		end_block_hdr = (kmm_block_hdr_t *)
			((((uint32_t)base_ptr) + initial_size) -(uint32_t) sizeof(kmm_block_hdr_t));
		//tprintf("   __end_block_hdr   = 0x%x\n",(uint32_t)end_block_hdr);
		memset((void *) end_block_hdr,0,sizeof(kmm_block_hdr_t));
		end_block_hdr->next_ptr = (uint32_t) end_block_hdr;
		end_block_hdr->prev_ptr = (uint32_t) start_block_hdr;
		end_block_hdr->flags = KMM_EOM_FLAG;
		start_block_hdr->next_ptr = (uint32_t) end_block_hdr;
		start_block_hdr->prev_ptr = (uint32_t) start_block_hdr;
		start_block_hdr->phys = page;
		start_block_hdr->flags = KMM_FREE_FLAG;
	}
	current_heap = heap;
	//tprintf(" kmm_setup_heap: return\n");
}

void kmm_merge_block(kmm_block_hdr_t *block){
	kmm_block_hdr_t *block_b;						
	block_b = (kmm_block_hdr_t *) block->next_ptr;				
	block->next_ptr = block_b->next_ptr;					 //Bn->next = Bn+2
	((kmm_block_hdr_t *)block_b->next_ptr)->prev_ptr = (uint32_t) block;//Bn+2->prev = Bn
	block->flags |= block_b->flags;							 //Merge flags
}

void kmm_split_block(kmm_block_hdr_t *block,size_t size){
	kmm_block_hdr_t *block_b;
	size_t block_size;
	block_size = (size_t) 
		(block->next_ptr - (uint32_t)(block + sizeof(kmm_block_hdr_t)));
	if ((block_size - size) < (sizeof(kmm_block_hdr_t) + KMM_MIN_BLOCK_SIZE))
		return;							//It'd be pointless or impossible to create a block this small.
	block_b = (kmm_block_hdr_t *) (block + sizeof(kmm_block_hdr_t) + size);
	memset((void *) block_b,0,sizeof(kmm_block_hdr_t));
	block_b->phys = block->phys + (sizeof(kmm_block_hdr_t) + size);
	block_b->next_ptr = block->next_ptr;
	block_b->prev_ptr = (uint32_t) block;
	block->next_ptr = (uint32_t) block_b;
	((kmm_block_hdr_t *)block_b->next_ptr)->prev_ptr = (uint32_t) block_b;//Bn+2->prev = Bn+1
	block_b->flags = KMM_FREE_FLAG;
}

void kmm_expand(uint8_t flags,kmm_block_hdr_t *last_block,size_t size){
	kmm_block_hdr_t *end_block_hdr;
	end_block_hdr = (kmm_block_hdr_t *)
		((((uint32_t)last_block) + size) - (uint32_t) sizeof(kmm_block_hdr_t));
	vmm_kbrk((void *)(last_block+sizeof(kmm_block_hdr_t)),size + KMM_GROW_SIZE);
	dprintf("kbrk done\n");
	memset((void *) end_block_hdr,0,sizeof(kmm_block_hdr_t));
	end_block_hdr->next_ptr = (uint32_t) end_block_hdr;
	end_block_hdr->prev_ptr = (uint32_t) last_block;
	end_block_hdr->flags = KMM_EOM_FLAG;
	last_block->next_ptr = (uint32_t) end_block_hdr;
	last_block->prev_ptr = (uint32_t) last_block;
	last_block->flags = KMM_FREE_FLAG | flags;
}

void *kmalloc(size_t size){
	kmm_block_hdr_t *current_block;
	size_t block_size;
	uint32_t p = palloc;
	palloc+=size;
	if (1==1)
		return (void *)p;

	//tprintf("[0x%X] - _kmalloc(0x%x \n",(void *)kmalloc,size);
	current_block = (kmm_block_hdr_t *) current_heap->free_ptr;	
kmalloc_loop:
	dtprintf("      _current_block = 0x%x\n",current_block);
	if ((current_block->flags & KMM_EOM_FLAG) != 0){//We've reached end of memory
		dprintf("      MSG: BRK\n");
		kmm_expand(0,current_block,size);

	}
	block_size = (size_t) 
		(current_block->next_ptr - (uint32_t)(current_block + sizeof(kmm_block_hdr_t)));
	dtprintf("      _block_size    = 0x%x\n",block_size);
	if ((current_block->flags & KMM_FREE_FLAG) != 0){//Free block :D
		dtprintf("      print: free block found\n");
		if (block_size < size)						 //Block too small :(
			goto kmalloc_recurse;
		if (block_size > size)       				 //Found large enough block but its too large
			kmm_split_block(current_block,size);
		dtprintf("      _current_block = 0x%x\n",current_block);
		current_block->flags &= ~KMM_FREE_FLAG;
	//	tprintf(" kmalloc:return 0x%x\n",(uint32_t) (((uint32_t)current_block) + (uint32_t)sizeof(kmm_block_hdr_t)));
		return (void *) (((uint32_t)current_block) + (uint32_t)sizeof(kmm_block_hdr_t));		
	}
kmalloc_recurse:
	current_block = (kmm_block_hdr_t *) current_block->next_ptr;
	goto kmalloc_loop;
}



void *kmalloc_ap(size_t size,physical_addr *phys){

	kmm_block_hdr_t *current_block;
	size_t block_size;
	uint32_t block_addr;
	//tprintf("[0x%X] - _kmalloc_ap(0x%x,0x%x) \n",(void *)kmalloc_ap,size,phys);
	current_block = (kmm_block_hdr_t *) current_heap->free_ptr;	 
kmallocap_loop:
	dtprintf("      _current_block = 0x%x\n",current_block);
	if ((current_block->flags & KMM_EOM_FLAG) != 0){//We've reached end of memory
		//TODO: When paging is done allocate new memory for heap
		//      Size: size + KMM_GROW_SIZE
		//DIRTYHACK: For now, return NULL pointer
		dprintf("      ERROR: Out of memory\n");
		return (void *) 0;
	}
	block_size = (size_t) 
		(current_block->next_ptr - (uint32_t)(current_block + sizeof(kmm_block_hdr_t)));
	dtprintf("      _block_size    = 0x%x\n",block_size);
	block_addr = (uint32_t) (current_block + sizeof(kmm_block_hdr_t));
	dtprintf("      _block_addr    = 0x%x\n",block_addr);
	if ((current_block->flags & KMM_FREE_FLAG) != 0){//Free block :D
		dtprintf("      print: free block found\n");
		if (block_size < size)						 //Block too small :(
			goto kmallocap_recurse;
		if ((block_addr & 0xFFF) != 0){//Block is not alligned
			dprintf("      print: block not alligned\n");
			uint32_t split_point = (0xFFF - (block_addr & 0xFFF)) - sizeof(kmm_block_hdr_t); ;
			dprintf ("0x%x 0x%x 0x%x 0x%x",block_addr,split_point,sizeof(kmm_block_hdr_t),block_size);
			if (block_size
				< split_point + 0x1000)
				goto kmallocap_recurse;
			kmm_split_block(current_block,split_point);
			//DIRTYHACK: We abuse kmm_split_block() to create a new block that starts at the 
			//DIRTYHACK: page boundary.
			current_block = (kmm_block_hdr_t *) current_block->next_ptr;
			dprintf("     ._current_block = 0x%x\n",current_block);
			block_size = (size_t) 
				(current_block->next_ptr - (uint32_t)(current_block + sizeof(kmm_block_hdr_t)));
			dprintf("     ._block_size    = 0x%x\n",block_size);
			block_addr = (uint32_t) (current_block + sizeof(kmm_block_hdr_t));
			dprintf("     ._block_addr    = 0x%x\n",block_addr);
		}

		if (block_size > size)       				 //Found large enough block but its too large
			kmm_split_block(current_block,size);
		current_block->flags &= ~KMM_FREE_FLAG;
		dprintf(" kmallocap:return 0x%x\n",(uint32_t) (((uint32_t)current_block) + (uint32_t)sizeof(kmm_block_hdr_t)));
		*phys = (physical_addr) (block_addr - 0xBFF00000);//Only works for remapped kheap;
		return (void *) block_addr;		
	}
kmallocap_recurse:
	current_block = (kmm_block_hdr_t *) current_block->next_ptr;
	goto kmallocap_loop;
}



void kfree(void *ptr){
	kmm_block_hdr_t *current_block;
	kmm_block_hdr_t *merge_block;

	if(1==1)
		return;

	//tprintf("[0x%X] - _kfree(0x%x) \n",(void *)kmalloc_ap,ptr);
	current_block = (kmm_block_hdr_t *) (((uint32_t)ptr) - (uint32_t)sizeof(kmm_block_hdr_t));
	dtprintf("      _current_block = 0x%x\n",current_block);
	if (current_block->prev_ptr != (uint32_t) current_block){
		merge_block = (kmm_block_hdr_t *) current_block->prev_ptr;
		dtprintf("      _merge_block   = 0x%x\n",merge_block);
		if ((merge_block->flags & KMM_FREE_FLAG) != 0){
			kmm_merge_block(merge_block);
			current_block = merge_block;
		}
	}
	merge_block = (kmm_block_hdr_t *) current_block->next_ptr;
	if ((merge_block->flags & KMM_FREE_FLAG) != 0)
		kmm_merge_block(current_block);
	current_block->flags |= KMM_FREE_FLAG;
	if (current_heap->free_ptr > (uint32_t) current_block)
		current_heap->free_ptr = (uint32_t) current_block;
}
